import os
import json
from typing import List, Dict, Optional
from langchain_deepseek import chat_models
from paragraph_write_agents.uilts.paragraph_state import ParagraphState


def _get_provider_conf_from_env(provider: str) -> Optional[Dict[str, str]]:
    key = provider.upper()
    api_key = os.getenv(f"{key}_API_KEY")
    base_url = os.getenv(f"{key}_BASE_URL")
    model = os.getenv(f"{key}_MODEL")
    if api_key and base_url and model:
        return {"api_key": api_key, "base_url": base_url, "model": model}
    return None


def get_llm(provider):
    conf = _get_provider_conf_from_env(provider)
    if conf:
        return chat_models.ChatDeepSeek(model=conf["model"], base_url=conf["base_url"], api_key=conf["api_key"],
                                        model_kwargs={"response_format": {"type": "json_object"}}
                                        )
    return None


def paragraph_select(state: ParagraphState) -> ParagraphState:
    """根据段落大纲选择分析类型"""
    outline = state.get("paragraph_outline")
    provider = state.get("provider")
    system_prompt = "你是一个科学计量报告大纲的专业分类人员，你需要理解用户的段落大纲，并根据段落大纲的内容，来判断用户需要撰写的段落类型。"
    user_prompt = f"""
        请阅读用户进行的输入，抽取并返回一个JSON对象，包含字段：
        - paragraph_diagram：用户需求呈现的分析范式，返回一个字符串，字符串只能是列表['文本型段落','数据型段落','总结型段落','其他']中的一个
        大纲：{outline}
        注意：
        - 判断依据：
            -如果用户的大纲要撰写的是关于引言、背景、意义、方法论等相关章节，是可以使用大模型直接生成的，那么需要撰写的段落属于'文本型段落'
            -如果用户的大纲要撰写的是对于一组数据、指标等进行的分析和解析，那么需要撰写的段落属于'数据型段落'
            -如果用户的大纲要撰写的内容是对前文中的总结，那么需要撰写的段落属于'总结型段落'
            -如果用户的大纲不属于以上三种类型，那么返回'其他'
        - 仅返回可解析的JSON字符串，不要包含解释或额外文本。
        """

    llm = get_llm(provider)
    msg = llm.invoke(
        [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    )
    parsed = json.loads(msg.content.strip())
    state["messages"].append(msg)
    state.update(parsed)
    return state


def fill_text_paragraph(state: ParagraphState) -> ParagraphState:
    """文本型段落"""
    outline = state.get("paragraph_outline")
    provider = state.get("provider")
    system_prompt = "你是一位擅长撰写科学计量学报告的资深研究员。请根据用户提供的需求，撰写一份专业、严谨、逻辑清晰的研究报告。"
    user_prompt = f"""
        用户的输入信息包括：
        - 引言大纲：{outline}
        - 具体技术/行业领域：{state.get("domain")}
        - 受众：{state.get("audience")}
        - 用户初始输入：{state.get("user_input")}
                
        输出的格式请返回一个JSON对象，包含字段：
        - paragraph_content：引言内容，返回一个字符串
        注意：
        - 仅返回可解析的JSON字符串，不要包含解释或额外文本。
        """

    llm = get_llm(provider)
    msg = llm.invoke(
        [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    )
    parsed = json.loads(msg.content.strip())
    print(parsed)
    state["messages"].append(msg)
    state.update(parsed)
    return state


def fill_data_paragraph(state: ParagraphState) -> ParagraphState:
    """填写数据段落"""
    print("fill_data_paragraph")
    provider = state.get("provider")
    outline = state.get("paragraph_outline")

    system_prompt = "你是一位擅长撰写科学计量学报告数据分析部分的资深研究员。请根据用户提供的结构化信息，为我撰写一份专业、严谨、逻辑清晰的研究报告数据分析段落。"
    user_prompt = f"""        
        用户的输入信息包括：
        - 数据分析大纲：{outline}
        - 具体技术/行业领域：{state.get("domain", "人工智能")}
        - 受众：{state.get("audience", "政府决策者")}
        - 用户初始输入：{state.get("user_input", "")}
        - 分析数据：{state.get("paragraph_data", "2024年学术论文、专利申请和产业动态数据")}
        - 分析方法：{state.get("analytical_method", "时间序列分析和主题建模")}
        - 图表类型：{state.get("chart_type", "折线图、柱状图")}
                
        输出的格式请返回一个JSON对象，包含字段：
        - paragraph_content：数据分析段落内容，返回一个字符串
        注意：
        - 仅返回可解析的JSON字符串，不要包含解释或额外文本。
        """

    llm = get_llm(provider)
    msg = llm.invoke(
        [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    )

    parsed = json.loads(msg.content.strip())
    print(parsed.get("paragraph_content", ""))

    state["messages"].append(msg)
    state.update(parsed)
    return state


def fill_summary_datagraph(state: ParagraphState) -> ParagraphState:
    """根据段落进行总结"""
    provider = state.get("provider")
    outline = state.get("paragraph_outline")

    system_prompt = "你是一位擅长撰写科学计量学报告结论和总结的资深研究员。请根据用户提供的结构化信息，为我撰写一份专业、严谨、逻辑清晰的研究报告总结段落。"
    user_prompt = f"""        
        用户的输入信息包括：
        - 总结大纲：{outline}
        - 具体技术/行业领域：{state.get("domain", "人工智能")}
        - 受众：{state.get("audience", "政府决策者")}
        - 用户初始输入：{state.get("user_input", "")}
                
        输出的格式请返回一个JSON对象，包含字段：
        - paragraph_content：总结段落内容，返回一个字符串
        注意：
        - 仅返回可解析的JSON字符串，不要包含解释或额外文本。
        """

    llm = get_llm(provider)
    msg = llm.invoke(
        [
            {"role": "system", "content": system_prompt},
            {"role": "user", "content": user_prompt},
        ]
    )
    parsed = json.loads(msg.content.strip())
    print(parsed.get("paragraph_content", ""))
    state["messages"].append(msg)
    state.update(parsed)
    return state


def route_paragraph(state: ParagraphState) -> str:
    # 根据 paragraph_outline 路由到对应节点
    val = (state.get("paragraph_diagram") or "").strip()
    mapping = {
        "文本型段落": "fill_text_paragraph",
        "数据型段落": "fill_data_paragraph",
        "总结型段落": "fill_summary_paragraph",
        "其他": "其他"
    }
    return mapping.get(val, "其他")
